package com.itheima.consultant;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

@Component
public class ZhiPuAiChatModel {

    @Value("${zhipu.api-key:}")
    private String apiKey;

    private final WebClient webClient;

    public ZhiPuAiChatModel(WebClient.Builder webClientBuilder) {
        this.webClient = webClientBuilder
                .baseUrl("https://open.bigmodel.cn/api/paas/v4")
                .build();
    }

    // 同步调用方法
    public String call(String message) {
        try {
            Map<String, Object> request = Map.of(
                    "model", "glm-4",
                    "messages", List.of(Map.of("role", "user", "content", message)),
                    "temperature", 0.4,
                    "max_tokens", 200
            );

            Map response = webClient.post()
                    .uri("/chat/completions")
                    .header("Authorization", "Bearer " + apiKey)
                    .header("Content-Type", "application/json")
                    .bodyValue(request)
                    .retrieve()
                    .bodyToMono(Map.class)
                    .block();
            System.out.println( response);
            System.out.println(response.get("choices"));
            // 提取回复内容
            List<Map<String, Object>> choices = (List<Map<String, Object>>) response.get("choices");
            if (choices != null && !choices.isEmpty()) {
                Map<String, Object> choice = choices.get(0);
                Map<String, Object> messageMap = (Map<String, Object>) choice.get("message");
                return (String) messageMap.get("content");
            } else {
                return "No choices found in the response.";
            }
        } catch (Exception e) {
            return "Error calling AI: " + e.getMessage();
        }
    }

    // 流式调用
    public Flux<String> stream(String message) {
        System.out.println("Stream request received with message: " + message);
        Map<String, Object> request = Map.of(
                "model", "glm-4",
                "messages", List.of(Map.of("role", "user", "content", message)),
                "temperature", 0.4,
                "max_tokens", 200,
                "stream", true
        );

        return webClient.post()
                .uri("/chat/completions")
                .header("Authorization", "Bearer " + apiKey)
                .header("Content-Type", "application/json")
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class)
                .doOnNext(sseData -> System.out.println("[SSE RAW]: " + sseData))
                .handle((sseData, sink) -> {
                    String content = extractContentFromSSE(sseData);
                    if (content != null && !content.trim().isEmpty()) {
                        System.out.println("[EMIT]: " + content);
                        sink.next(content); // 发射纯文本内容
                    }
                });

    }

    private String extractContentFromSSE(String sseData) {
        if (sseData == null || !sseData.startsWith("data:")) {
            return null;
        }

        try {
            String jsonPart = sseData.substring(5).trim(); // 去掉 "data: "
            if ("[DONE]".equals(jsonPart)) {
                return null; // 流结束
            }

            ObjectMapper objectMapper = new ObjectMapper();
            Map<String, Object> response = objectMapper.readValue(jsonPart, Map.class);
            List<Map<String, Object>> choices = (List<Map<String, Object>>) response.get("choices");

            if (choices != null && !choices.isEmpty()) {
                Map<String, Object> choice = choices.get(0);
                Map<String, Object> delta = (Map<String, Object>) choice.get("delta");

                if (delta != null) {
                    Object contentObj = delta.get("content");
                    if (contentObj instanceof String str && !str.trim().isEmpty()) {
                        return str;
                    }
                }
            }
        } catch (Exception e) {
            // 忽略解析失败的块（如心跳、错误等）
            return null;
        }
        return null;
    }
}